/*
 * HSM Proxy Project.
 * Copyright (C) 2013 FedICT.
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License version
 * 3.0 as published by the Free Software Foundation.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, see 
 * http://www.gnu.org/licenses/.
 */

package be.fedict.hsm.model;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.security.KeyStore;
import java.security.KeyStore.PrivateKeyEntry;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.Security;
import java.security.UnrecoverableEntryException;
import java.security.cert.CertificateException;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;

import javax.ejb.Local;
import javax.ejb.Stateless;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import sun.security.pkcs11.SunPKCS11;
import be.fedict.hsm.entity.KeyStoreEntity;
import be.fedict.hsm.entity.KeyStoreState;
import be.fedict.hsm.entity.KeyStoreType;

@Stateless
@Local(KeyStoreLoader.class)
public class KeyStoreLoaderBean implements KeyStoreLoader {

	private static final Log LOG = LogFactory.getLog(KeyStoreLoaderBean.class);

	@Override
	public Map<String, PrivateKeyEntry> loadKeyStore(
			KeyStoreEntity keyStoreEntity) {
		KeyStoreType keyStoreType = keyStoreEntity.getKeyStoreType();
		Map<String, PrivateKeyEntry> result;
		switch (keyStoreType) {
		case PKCS12:
			try {
				result = loadPKCS12(keyStoreEntity);
			} catch (Exception e) {
				LOG.error("error loading PKCS#12 key store: " + e.getMessage(),
						e);
				keyStoreEntity.setKeyStoreState(KeyStoreState.ERROR);
				return null;
			}
			break;
		case PKCS11:
			try {
				result = loadPKCS11(keyStoreEntity);
			} catch (Exception e) {
				LOG.error("error loading PKCS#11 key store: " + e.getMessage(),
						e);
				keyStoreEntity.setKeyStoreState(KeyStoreState.ERROR);
				return null;
			}
			break;
		default:
			throw new IllegalArgumentException("unsupported key store type: "
					+ keyStoreType);
		}
		keyStoreEntity.setKeyStoreState(KeyStoreState.ACTIVE);
		return result;
	}

	private Map<String, PrivateKeyEntry> loadPKCS11(
			KeyStoreEntity keyStoreEntity) throws KeyStoreException,
			IOException, NoSuchAlgorithmException, CertificateException,
			UnrecoverableEntryException {
		File tmpConfigFile = File.createTempFile("pkcs11-", ".conf");
		tmpConfigFile.deleteOnExit();
		PrintWriter configWriter = new PrintWriter(new FileOutputStream(
				tmpConfigFile));
		configWriter.println("name=HSM-" + keyStoreEntity.getId());
		String path = keyStoreEntity.getPath();
		LOG.debug("PKCS11 path: " + path);
		LOG.debug("slot list index: " + keyStoreEntity.getSlotListIndex());
		configWriter.println("library=" + path);
		configWriter.println("slotListIndex="
				+ keyStoreEntity.getSlotListIndex());
		configWriter.close();
		SunPKCS11 sunPKCS11 = new SunPKCS11(tmpConfigFile.getAbsolutePath());
		LOG.debug("adding SunPKCS11 JCA provider: " + sunPKCS11.getName());
		/*
		 * Reloads also need to work properly.
		 */
		Security.removeProvider(sunPKCS11.getName());
		Security.addProvider(sunPKCS11);
		KeyStore keyStore = KeyStore.getInstance("PKCS11", sunPKCS11);
		if (null != keyStoreEntity.getPassword()) {
			keyStore.load(null, keyStoreEntity.getPassword().toCharArray());
		} else {
			keyStore.load(null, null);
		}
		String keyStorePassword = keyStoreEntity.getPassword();
		return loadKeys(keyStoreEntity, keyStore, keyStorePassword);
	}

	private Map<String, PrivateKeyEntry> loadPKCS12(
			KeyStoreEntity keyStoreEntity) throws KeyStoreException,
			NoSuchAlgorithmException, CertificateException, IOException,
			UnrecoverableEntryException {
		String keyStorePath = keyStoreEntity.getPath();
		InputStream keyStoreInputStream = new FileInputStream(keyStorePath);
		KeyStore keyStore = KeyStore.getInstance("PKCS12");
		String keyStorePassword = keyStoreEntity.getPassword();
		keyStore.load(keyStoreInputStream, keyStorePassword.toCharArray());
		return loadKeys(keyStoreEntity, keyStore, keyStorePassword);
	}

	private Map<String, PrivateKeyEntry> loadKeys(
			KeyStoreEntity keyStoreEntity, KeyStore keyStore,
			String keyStorePassword) throws KeyStoreException,
			NoSuchAlgorithmException, UnrecoverableEntryException {
		Enumeration<String> aliases = keyStore.aliases();
		Map<String, PrivateKeyEntry> keyStorePrivateKeys = new HashMap<String, PrivateKeyEntry>();
		while (aliases.hasMoreElements()) {
			String alias = aliases.nextElement();
			PrivateKeyEntry privateKeyEntry;
			if (null != keyStorePassword) {
				privateKeyEntry = (PrivateKeyEntry) keyStore.getEntry(
						alias,
						new KeyStore.PasswordProtection(keyStorePassword
								.toCharArray()));
			} else {
				privateKeyEntry = (PrivateKeyEntry) keyStore.getEntry(alias,
						null);
			}
			keyStorePrivateKeys.put(alias, privateKeyEntry);
		}
		return keyStorePrivateKeys;
	}
}
